using LinearAlgebra

"""
    stochastic_block_model(c, n)

Return a Graph generated according to the Stochastic Block Model (SBM).

`c[a,b]` : Mean number of neighbors of a vertex in block `a` belonging to block `b`.
           Only the upper triangular part is considered, since the lower traingular is
           determined by ``c[b,a] = c[a,b] * \\frac{n[a]}{n[b]}``.
`n[a]` : Number of vertices in block `a`

### Optional Arguments
- `seed=-1`: set the RNG seed.

For a dynamic version of the SBM see the [`StochasticBlockModel`](@ref) type and
related functions.
"""
function stochastic_block_model(c::Matrix{T}, n::Vector{U}; seed::Int=-1) where T <: Real where U <: Integer
    size(c, 1) == size(c, 2) == length(n) || throw(ArgumentError("matrix-vector size mismatch"))

    # init dsfmt generator without altering GLOBAL_RNG
    rng = getRNG(seed)
    N = sum(n)
    K = length(n)
    nedg = zeros(Int, K, K)
    g = SimpleGraph(N)
    cum = [sum(n[1:a]) for a = 0:K]
    for a = 1:K
        ra = (cum[a] + 1):cum[a + 1]
        for b = a:K
            ((a == b) && !(c[a, b] <= n[b] - 1)) || ((a != b) && !(c[a, b] <= n[b])) &&
                error("Mean degree cannot be greater than available neighbors in the block.") # TODO 0.7: turn into some other error?

            m = a == b ? div(n[a] * (n[a] - 1), 2) : n[a] * n[b]
            p = a == b ? n[a] * c[a, b] / (2m) : n[a] * c[a, b] / m
            nedg = randbn(m, p, seed)
            rb = (cum[b] + 1):cum[b + 1]
            i = 0
            while i < nedg
                source = rand(rng, ra)
                dest = rand(rng, rb)
                if source != dest
                    if add_edge!(g, source, dest)
                        i += 1
                    end
                end
            end
        end
    end
    return g
end

"""
    stochastic_block_model(cint, cext, n)

Return a Graph generated according to the Stochastic Block Model (SBM), sampling
from an SBM with ``c_{a,a}=cint``, and ``c_{a,b}=cext``.
"""
function stochastic_block_model(cint::T, cext::T, n::Vector{U}; seed::Int=-1) where T <: Real where U <: Integer
    K = length(n)
    c = [ifelse(a == b, cint, cext) for a = 1:K, b = 1:K]
    stochastic_block_model(c, n, seed=seed)
end

"""
    StochasticBlockModel{T,P}

A type capturing the parameters of the SBM.
Each vertex is assigned to a block and the probability of edge `(i,j)`
depends only on the block labels of vertex `i` and vertex `j`.

The assignement is stored in nodemap and the block affinities a `k` by `k`
matrix is stored in affinities.

`affinities[k,l]` is the probability of an edge between any vertex in
block `k` and any vertex in block `l`.

### Implementation Notes
Graphs are generated by taking random ``i,j ∈ V`` and
flipping a coin with probability `affinities[nodemap[i],nodemap[j]]`.
"""
mutable struct StochasticBlockModel{T <: Integer,P <: Real}
    n::T
    nodemap::Array{T}
    affinities::Matrix{P}
end

==(sbm::StochasticBlockModel, other::StochasticBlockModel) =
    (sbm.n == other.n) && (sbm.nodemap == other.nodemap) && (sbm.affinities == other.affinities)


# A constructor for StochasticBlockModel that uses the sizes of the blocks
# and the affinity matrix. This construction implies that consecutive
# vertices will be in the same blocks, except for the block boundaries.
function StochasticBlockModel(sizes::AbstractVector, affinities::AbstractMatrix)
    csum = cumsum(sizes)
    j = 1
    nodemap = zeros(Int, csum[end])
    for i in 1:csum[end]
        if i > csum[j]
            j += 1
        end
        nodemap[i] = j
    end
    return StochasticBlockModel(csum[end], nodemap, affinities)
end


### TODO: This documentation needs work. sbromberger 20170326
"""
    sbmaffinity(internalp, externalp, sizes)

Produce the sbm affinity matrix with internal probabilities `internalp`
and external probabilities `externalp`.
"""
function sbmaffinity(internalp::Vector{T}, externalp::Real, sizes::Vector{U}) where T <: Real where U <: Integer
    numblocks = length(sizes)
    numblocks == length(internalp) || throw(ArgumentError("Inconsistent input dimensions: internalp, sizes"))
    B = diagm(0 => internalp) + externalp * (ones(numblocks, numblocks) - I)
    return B
end

function StochasticBlockModel(internalp::Real,
                              externalp::Real,
                              size::Integer,
                              numblocks::Integer)
    sizes = fill(size, numblocks)
    B = sbmaffinity(fill(internalp, numblocks), externalp, sizes)
    StochasticBlockModel(sizes, B)
end

function StochasticBlockModel(internalp::Vector{T}, externalp::Real, sizes::Vector{U}) where {T <: Real, U <: Integer}
    B = sbmaffinity(internalp, externalp, sizes)
    return StochasticBlockModel(sizes, B)
end


const biclique = ones(2, 2) - Matrix{Float64}(I, 2, 2)

#TODO: this documentation needs work. sbromberger 20170326
"""
    nearbipartiteaffinity(sizes, between, intra)

Construct the affinity matrix for a near bipartite SBM.
`between` is the affinity between the two parts of each bipartite community.
`intra` is the probability of an edge within the parts of the partitions.

This is a specific type of SBM with ``\\frac{k}{2} blocks each with two halves.
Each half is connected as a random bipartite graph with probability `intra`
The blocks are connected with probability `between`.
"""
function nearbipartiteaffinity(sizes::AbstractVector{T}, between::Real, intra::Real) where {T <: Integer}
    numblocks = div(length(sizes), 2)
    return kron(between * Matrix{Float64}(I, numblocks, numblocks), biclique) + Matrix{Float64}(I, 2 * numblocks, 2 * numblocks) * intra
end

#Return a generator for edges from a stochastic block model near-bipartite graph.
nearbipartiteaffinity(sizes::Vector{T}, between::Real, inter::Real, noise::Real) where {T <: Integer} =
    nearbipartiteaffinity(sizes, between, inter) .+ noise

nearbipartiteSBM(sizes, between, inter, noise) =
    StochasticBlockModel(sizes, nearbipartiteaffinity(sizes, between, inter, noise))

"""
    random_pair(rng, n)

Generate a stream of random pairs in `1:n` using random number generator `RNG`.
"""
function random_pair(rng::AbstractRNG, n::Integer)
    f(ch) = begin
        while true
            put!(ch, SimpleEdge(rand(rng, 1:n), rand(rng, 1:n)))
        end
    end
    return f
end


"""
    make_edgestream(sbm)

Take an infinite sample from the Stochastic Block Model `sbm`.
Pass to `Graph(nvg, neg, edgestream)` to get a Graph object based on `sbm`.
"""
function make_edgestream(sbm::StochasticBlockModel, rng::AbstractRNG = GLOBAL_RNG)
    pairs = Channel(random_pair(rng, sbm.n), ctype=SimpleEdge, csize=32)
    edges(ch) = begin
        for e in pairs
            i, j = Tuple(e)
            i == j && continue
            p = sbm.affinities[sbm.nodemap[i], sbm.nodemap[j]]
            if rand(rng) < p
                put!(ch, e)
            end
        end
    end
    return Channel(edges, ctype=SimpleEdge, csize=32)
end

"""
    SimpleGraph{T}(nv, ne, edgestream::Channel)

Construct a `SimpleGraph{T}` with `nv` vertices and `ne` edges from `edgestream`.
Can result in less than `ne` edges if the channel `edgestream` is closed prematurely.
Duplicate edges are only counted once.
The element type is the type of `nv`.
"""
function SimpleGraph(nvg::Integer, neg::Integer, edgestream::Channel)
    g = SimpleGraph(nvg)
    # println(g)
    for e in edgestream
        add_edge!(g, e)
        ne(g) >= neg && break
    end
    return g
end

"""
    SimpleGraph{T}(nv, ne, smb::StochasticBlockModel)

Construct a random `SimpleGraph{T}` with `nv` vertices and `ne` edges.
The graph is sampled according to the stochastic block model `smb`.
The element type is the type of `nv`.
"""
SimpleGraph(nvg::Integer, neg::Integer, sbm::StochasticBlockModel) =
    SimpleGraph(nvg, neg, make_edgestream(sbm))

#TODO: this documentation needs work. sbromberger 20170326
"""
    blockcounts(sbm, A)

Count the number of edges that go between each block.
"""
function blockcounts(sbm::StochasticBlockModel, A::AbstractMatrix)
    I = collect(1:sbm.n)
    J =  [sbm.nodemap[i] for i in 1:sbm.n]
    V =  ones(sbm.n)
    Q = sparse(I, J, V)
    return (Q'A) * Q
end


function blockcounts(sbm::StochasticBlockModel, g::AbstractGraph)
    return blockcounts(sbm, adjacency_matrix(g))
end

function blockfractions(sbm::StochasticBlockModel, g::Union{AbstractGraph,AbstractMatrix})
    bc = blockcounts(sbm, g)
    bp = bc ./ sum(bc)
    return bp
end

